﻿#pragma once

#include <assert.h>

namespace dr
{
	template<class K, class V>
	struct AVLTreeNode
	{
		// 需要parent指针，后续更新平衡因⼦可以看到
		pair<K, V> _kv;
		AVLTreeNode<K, V>* _left;
		AVLTreeNode<K, V>* _right;
		AVLTreeNode<K, V>* _parent;
		int _bf; // balance factor
		AVLTreeNode(const pair<K, V>& kv)
			:_kv(kv)
			, _left(nullptr)
			, _right(nullptr)
			, _parent(nullptr)
			, _bf(0)
		{}
	};

	template<class K, class V>
	class AVLTree
	{
		typedef AVLTreeNode<K, V> Node;
	public:
		//...
		bool insert(const pair<K, V>& kv)
		{
			if (_root == nullptr)
			{
				_root = new Node(kv);
			}
			else
			{
				Node* cur = _root;
				Node* parent = nullptr;
				while (cur)
				{
					if (kv.first < cur->_kv.first)
					{
						parent = cur;
						cur = cur->_left;
					}
					else if (kv.first > cur->_kv.first)
					{
						parent = cur;
						cur = cur->_right;
					}
					else
					{
						return false;
					}
				}

				cur = new Node(kv);

				if (kv.first < parent->_kv.first)
				{
					parent->_left = cur;
				}
				else
				{
					parent->_right = cur;
				}
				cur->_parent = parent;

				// 更新平衡因子
				while (parent)
				{
					if (cur == parent->_left)
					{
						parent->_bf--;
					}
					else if (cur == parent->_right)
					{
						parent->_bf++;
					}

					if (parent->_bf == 0)
					{
						break;
					}
					else if (parent->_bf == -1 || parent->_bf == 1)
					{
						cur = parent;
						parent = cur->_parent;
					}
					else if (parent->_bf == -2 || parent->_bf == 2) // 旋转
					{
						if (parent->_bf == -2 && cur->_bf == -1)
						{
							// 右单旋
							RotateR(parent);
						}
						else if (parent->_bf == 2 && cur->_bf == 1)
						{
							// 左单旋
							RotateL(parent);
						}
						else if (parent->_bf == -2 && cur->_bf == 1)
						{
							RotateLR(parent);
						}
						else if (parent->_bf == 2 && cur->_bf == -1)
						{
							RotateRL(parent);
						}
						break;
					}
					else
					{
						assert(false);
					}
				}
			}
			return true;
		}

		// 查找
		Node* Find(const K& key)
		{
			Node* cur = _root;
			while (cur)
			{
			if (cur->_kv.first < key)
			{
			cur = cur->_right;
			}
			else if (cur->_kv.first > key)
			{
			cur = cur->_left;
			}
			else
			{
			return cur;
			}
			}
			return nullptr;
		}

		void InOrder()
		{
			_InOrder(_root);
			cout << endl;
		}

		bool IsBalanceTree()
		{
			return _IsBalanceTree(_root);
		}

		int Height()
		{
			return _Height(_root);
		}

		int Size()
		{
			return _Size(_root);
		}

	private:

		void _InOrder(Node* root)
		{
			if (root == nullptr)
			{
				return;
			}
			else
			{
				_InOrder(root->_left);
				cout << root->_kv.first << " " << root->_kv.second << endl;
				_InOrder(root->_right);
			}
		}

		// 平衡检测
		int _Height(Node* root)
		{
			if (root == nullptr)
				return 0;
			int leftHeight = _Height(root->_left);
			int rightHeight = _Height(root->_right);
			return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
		}

		bool _IsBalanceTree(Node* root)
		{
			// 空树也是AVL树
			if (nullptr == root)
				return true;
			// 计算pRoot结点的平衡因⼦：即pRoot左右⼦树的⾼度差
			int leftHeight = _Height(root->_left);
			int rightHeight = _Height(root->_right);
			int diff = rightHeight - leftHeight;
			// 如果计算出的平衡因⼦与pRoot的平衡因⼦不相等，或者
			// pRoot平衡因⼦的绝对值超过1，则⼀定不是AVL树
			if (abs(diff) >= 2)
			{
				cout << root->_kv.first << "⾼度差异常" << endl;
				return false;
			}
			if (root->_bf != diff)
			{
				cout << root->_kv.first << "平衡因⼦异常" << endl;
				return false;
			}
			// pRoot的左和右如果都是AVL树，则该树⼀定是AVL树
			return _IsBalanceTree(root->_left) && _IsBalanceTree(root->_right);
		}

		int _Size(Node* root)
		{
			if (root == nullptr)
			{
				return 0;
			}

			return _Size(root->_left) + _Size(root->_right) + 1;
		}

		void RotateR(Node* parent)
		{
			Node* subL = parent->_left;
			Node* subLR = subL->_right;

			Node* ppNode = parent->_parent;

			parent->_left = subLR;
			if (subLR)
				subLR->_parent = parent;
			subL->_right = parent;
			parent->_parent = subL;

			if (ppNode == nullptr)
			{
				_root = subL;
				subL->_parent = nullptr;
			}
			else
			{
				if (ppNode->_left == parent)
				{
					ppNode->_left = subL;

				}
				else
				{
					ppNode->_right = subL;
				}
				subL->_parent = ppNode;
			}

			parent->_bf = subL->_bf = 0;
		}

		void RotateL(Node* parent)
		{
			Node* subR = parent->_right;
			Node* subRL = subR->_left;

			Node* ppNode = parent->_parent;

			parent->_right = subRL;
			if (subRL)
				subRL->_parent = parent;
			subR->_left = parent;
			parent->_parent = subR;

			if (ppNode == nullptr)
			{
				_root = subR;
				subR->_parent = nullptr;
			}
			else
			{
				if (ppNode->_left == parent)
				{
					ppNode->_left = subR;
				}
				else
				{
					ppNode->_right = subR;
				}
				subR->_parent = ppNode;
			}

			subR->_bf = parent->_bf = 0;
		}

		void RotateLR(Node* parent)
		{
			Node* kidL = parent->_left;
			Node* kidLR = kidL->_right;
			int bf = kidLR->_bf;

			RotateL(kidL);
			RotateR(parent);

			if (bf == 1)
			{
				parent->_bf = 0;
				kidLR->_bf = 0;
				kidL->_bf = -1;
			}
			else if (bf == -1)
			{
				parent->_bf = 1;
				kidLR->_bf = 0;
				kidL->_bf = 0;
			}
			else if (bf == 0)
			{
				parent->_bf = 0;
				kidLR->_bf = 0;
				kidL->_bf = 0;
			}
			else
			{
				assert(false);
			}
		}

		void RotateRL(Node* parent)
		{
			Node* kidR = parent->_right;
			Node* kidRL = kidR->_left;
			int bf = kidRL->_bf;

			RotateR(kidR);
			RotateL(parent);

			if (bf == 1)
			{
				parent->_bf = -1;
				kidRL->_bf = 0;
				kidR->_bf = 0;
			}
			else if (bf == -1)
			{
				parent->_bf = 0;
				kidRL->_bf = 0;
				kidR->_bf = 1;
			}
			else if (bf == 0)
			{
				parent->_bf = 0;
				kidRL->_bf = 0;
				kidR->_bf = 0;
			}
			else
			{
				assert(false);
			}
		}

	private:
		Node* _root = nullptr;
	};
}
